import cv2
import numpy as np
import os
from sam_predict import load_model, do_infer
import time
# 全局变量，用于保存鼠标点击的像素坐标
clicked_point = None

def mouse_callback(event, x, y, flags, param):
    """鼠标点击回调函数"""
    global clicked_point
    if event == cv2.EVENT_LBUTTONDOWN:  # 左键点击
        clicked_point = (x, y)
        print(f"选中的像素点坐标: {clicked_point}")



def get_image_paths(directory):
    # 定義常見圖片格式
    image_extensions = (".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff", ".webp")
    image_paths = []

    # 遍歷目錄及其子目錄
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.lower().endswith(image_extensions):  # 判斷是否是圖片
                image_paths.append(os.path.join(root, file))

    image_paths.sort()  # 按完整路徑排序
    return image_paths


def read_txt_to_2d_list(filename):
    """
    从文本文件中读取数据并存储为二维列表。

    参数:
        filename (str): 文本文件名。

    返回:
        list: 包含所有行数据的二维列表，每行数据是一个子列表。
    """
    # 存储结果的二维列表
    data_list = []

    try:
        # 打开并读取文件
        with open(filename, 'r') as file:
            for line in file:
                # 去除行末的换行符
                line = line.strip()
                # 分割行中的值，将字符串转为整数
                values = list(map(int, line.split(', ')))
                # 将分割的值作为子列表添加到二维列表中
                data_list.append(values)
    except FileNotFoundError:
        print(f"文件 '{filename}' 未找到，请检查文件路径！")
    except ValueError:
        print("文件内容格式错误，无法转换为数字！")
    
    return data_list



def is_in_range(value, list1):
    """
    判断一个数是否在指定区间内（包括断点）。

    参数:
        value (int): 要判断的数。
        start (int): 区间起始值，默认为 1066。
        end (int): 区间结束值，默认为 1079。

    返回:
        bool: 如果在区间内返回 True，否则返回 False。
    """
    start, end = list1

    return start <= value <= end

def main():
    global clicked_point  # 声明 clicked_point 为全局变量
    checkpoint_path = '/home/ai/wlm/gitee/deep-learing/models/sam_vit_h_4b8939.pth'
    image_path = "/home/JSDC/017254/code/gitee/deep-learing/imgs/1744185603703480.jpg"  # Replace with your actual image path
    predictor = load_model(checkpoint_path)
    # masks, scores, logits = do_infer(input_points, input_labels, predictor, image_path)
    # 使用示例
    txt_path = "/home/ai/wlm/gitee/deep-learing/txt_center/clicked_points.txt"

    directory_path = '/home/ai/wlm/gitee/deep-learing/imgs_src/left_undistort_fisheye'
    valide_values = read_txt_to_2d_list(txt_path)

    f1 = [396, 505]
    f2 = [612, 628]
    f3 = [829, 852]
    f4 = [937, 1059]
    f5 = [1066, 1079]

    images = get_image_paths(directory_path)
    for index, img in enumerate(images):  # 添加序号
        ind = index + 1
        print(f'start to deal {ind} data')
        if is_in_range(ind, f1) or is_in_range(ind,f2) or is_in_range(ind, f3) or is_in_range(ind,f4) or is_in_range(ind, f5):
            print(f'id: {index+1} continue')
            continue
        # print(img)
        # rgb_image = cv2.imread(img)
        else:
            id, file_id, u, v = valide_values[index]
            input_points = [[u,v]]
            input_labels = [1]

            masks, scores, logits = do_infer(input_points, input_labels, predictor, img)
            file_name_with_ext = os.path.basename(img)  # 提取 "1744185603703480.jpg"

            # 去掉擴展名
            file_id = os.path.splitext(file_name_with_ext)[0]  # 提取 "1744185603703480"
            # print(f'file_id: {file_id}')
            for i in range(1):
                mask = masks[i]  # 獲取每個遮罩 (形狀為 768x1024)
                
                # 將遮罩數據轉換為 0-255 的範圍（如果不是）
                mask = (mask * 255).astype(np.uint8)  # 確保數據是整數類型，範圍在 0-255
                
                # 保存遮罩為圖片
                output_filename = f"/home/ai/wlm/gitee/deep-learing/masks/mask_{file_id}.png"  # 生成文件名
                cv2.imwrite(output_filename, mask)  # 使用 OpenCV 保存圖片
                print(f"Saved mask {i} as {output_filename}")

if __name__ == "__main__":
    main()
